import tushare as ts
import pandas as pd
import time
import matplotlib.pyplot as plt
import os
from openpyxl import load_workbook
from openpyxl.drawing.image import Image
import xlwings as xw
import multiprocessing


os.chdir('/Users/liujiannan/Desktop/GP_PY')
token = '8b334e2d5d482292aa8b516c7f59dd6d23c0266dac3dc407dfb63ea1'
pro = ts.pro_api(token)


def tedliu(ts_code_p):
    df = pro.daily(ts_code=ts_code_p, start_date='20201201', end_date='20201229', fields='ts_code,trade_date,vol')
    df.replace("202012", "", regex=True, inplace=True)
    df = df.T
    df.columns = df.iloc[1]
    df.drop("trade_date", inplace=True)
    df.rename(index={"vol": ts_code_p}, inplace=True)
    df.drop("ts_code", inplace=True)
    a = df.columns.values.tolist()
    b = df.loc[ts_code_p]
    # plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.xlabel('Daily')
    plt.ylabel('Value')
    plt.plot(a, b)
    plt.savefig((str(ts_code_p) + '.png'), dpi=100)
    plt.close()
    return df


if __name__ == "__main__":
    print("...Starting...")
    now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    print(now)
    now1 = time.time()
    data = pro.query('stock_basic', exchange='', list_status='L', fields='ts_code,name')
    pool = multiprocessing.Pool(processes=6)
    multiple_results = [pool.apply_async(tedliu, args=(i,)) for i in data['ts_code']]
    # multiple_results = [pool.apply_async(tedliu, args=(i,)) for i in ['000001.SZ', '000002.SZ', '000004.SZ', '000005.SZ']]
    lst = [res.get() for res in multiple_results]
    pool.close()
    pool.join()

    result = pd.concat(lst)
    result.to_excel('phase1.xlsx')

    wb = load_workbook('phase1.xlsx')
    ws = wb.active

    for i in range(2, len(lst) + 2):
        ws.column_dimensions['Y'].width = 65
        ws.row_dimensions[i].height = 120
        img = Image(str(ws.cell(row=i, column=1).value) + '.png')  # 缩放图片
        img.width, img.height = (500, 154)
        ws.add_image(img, ("Y" + str(i)))  # 图片 插入 A1 的位置上

    wb.save('phase1.xlsx')  # 新的结果保存输出

    print("------------------------------------")
    print("...Finished...")
    now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    print(now)
    now2 = time.time()
    now3 = now2 - now1
    print("...用时{} Secs...".format(round(now3, 2)))
